{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jtLiiWFaZlUb"
   },
   "source": [
    "\n",
    "# **Paper Information**\n",
    "\n",
    "\n",
    "**TransGAN: Two Transformers Can Make One Strong GAN, and That Can Scale Up, CVPR 2021**, Yifan Jiang, Shiyu Chang, Zhangyang Wang\n",
    "\n",
    "*   Paper Link: https://arxiv.org/pdf/2102.07074v2.pdf\n",
    "*   Official Implementation: https://github.com/VITA-Group/TransGAN\n",
    "*   Paper Presentation by Ahmet Sarıgün : https://www.youtube.com/watch?v=xwrUkHiDoiY\n",
    "\n",
    "\n",
    "**Project Group Members:**\n",
    "\n",
    "\n",
    "*   Ahmet Sarıgün, ahmet.sarigun@metu.edu.tr\n",
    "*   Dursun Bekci, bekci.dursun@metu.edu.tr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GUv7BXE1ZrWc"
   },
   "source": [
    "## **Paper Summary**\n",
    "### **Introduction**\n",
    "TransGAN is a transformer-based GAN model which can be considered as a pilot study as being completely free of convolutions. The architecture of TransGAN mainly consists of a memory-friendly transformer-based generator that progressively increases feature resolution, and correspondingly a patch-level discriminator that is also transformer-based. In training of the model, a series of techniques are combined in the original paper such as data augmentation, modified normalization, and relative position encoding to overcome the general training instability issues of the GAN. We implemented data augmentation [(Dosovitskiy et al., 2020)](https://arxiv.org/pdf/2010.11929.pdf), and relative position encoding in our work. In the original paper, performance of the model tested on different datasets such as STL-10, CIFAR-10, CelebA datasets and achieved competitive results compared to current state-of-the-art GANs using convolutions. In our project, we only tested our implementation on CIFAR10 dataset as we stated in our experimental result goals.\n",
    "\n",
    "### **TransGAN Architecture**\n",
    "The architecture pipeline of TransGAN is shown below in the figure taken from the original paper.\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/images/transgan.jpg\"> \n",
    "    \n",
    "Figure 1: The pipeline of the pure transform-based generator and discriminator of TransGAN.\n",
    "\n",
    "### **Transformer Encoder as Basic Block**\n",
    "We used the transformer encoder [(Vaswani et al., 2017)](https://arxiv.org/pdf/1706.03762.pdf) as our basic block as in the original paper. An encoder is a composition of two parts. The first part is constructed by a multi-head self-attention module and the second part is a feed-forward MLP with GELU non-linearity. We apply layer normalization [(Ba et al., 2016)](https://arxiv.org/pdf/1607.06450.pdf) before both of the two parts.Both parts employ residual connection.\n",
    "\n",
    "$$Attention(Q, K, V ) = softmax(QK^T√d_k)V$$\n",
    "<img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/images/vit.gif\">\n",
    "Credits for illustration of ViT: [@lucidrains](https://github.com/lucidrains)\n",
    "\n",
    "\n",
    "### **Memory-friendly Generator**\n",
    "In building the memory-friendly generator, TransGAN utilizes a common design philosophy in CNN-based GANs which iteratively upscale the resolution at multiple stages. Figure 1 (left) illustrates the memory-friendly generator which consists of multiple stages with several transformer blocks. At each stage, feature map resolution is gradually increased until it meets the target resolution *H × W*. The generator takes the random noise input and passes it through a multiple-layer perceptron (MLP). The output vector reshaped into a $H_0 × W_0$ resolution feature map (by default $H_0$ = $W_0$ = 8), each point a C-dimensional embedding.  This “feature map\" is next treated as a length-64 sequence of C-dimensional tokens, combined with the learnable positional encoding. \n",
    "Then, transformer encoders take embedding tokens as inputs and calculate the correspondence between each token recursively. To synthesize higher resolution images, we insert an upsampling module after each stage, consisting of a pixelshuffle [(Shi et al., 2016)](https://arxiv.org/pdf/1609.05158.pdf) module. \n",
    "\n",
    "### **Tokenized-input for Discriminator**\n",
    "The authors design the discriminator as shown in Figure 1 (right) that it takes the patches of an image as inputs. Then, they split the input images $Y$ ∈ $R^{H×W×3}$ into 8x8 patches where each patch can be regarded as a \"word\". The patches are then converted to the 1D sequence of token embeddings through a linear flatten layer. After that, learnable position encoding is added, and tokens pass through the transformer encoder. Finally, tokens are taken by the classification head to output the real/fake prediction.    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hBtZpS3T-cl7"
   },
   "source": [
    "### **Training the Model**\n",
    "In this section, we show our training code and training score for CIFAR-10 Dataset with the best performance hyperparameters that we found. We trained the largest model TransGAN-XL with data augmentation using different hyperparameters, and record the results in cifar/experiments folder.    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "718NzrFLOUmr"
   },
   "source": [
    "## Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0GEJLtggOUms",
    "outputId": "6b9e1d8e-f221-4914-ef7e-1971eddc823e",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using downloaded and verified file: ./data\\cifar-10-python.tar.gz\n",
      "Extracting ./data\\cifar-10-python.tar.gz to ./data\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import time\n",
    "import argparse\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision.utils import make_grid, save_image\n",
    "\n",
    "from tensorboardX import SummaryWriter\n",
    "from tqdm import tqdm\n",
    "from copy import deepcopy\n",
    "\n",
    "\n",
    "from utils import *\n",
    "from models import *\n",
    "from fid_score import *\n",
    "from inception_score import *\n",
    "\n",
    "!mkdir checkpoint\n",
    "!mkdir generated_imgs\n",
    "!pip install tensorboardX\n",
    "!mkdir fid_stat\n",
    "%cd fid_stat\n",
    "!wget bioinf.jku.at/research/ttur/ttur_stats/fid_stats_cifar10_train.npz\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BNBeDCAXL6WI"
   },
   "source": [
    "## Hyperparameters for CIFAR-10 Dataset\n",
    "Since Google Colab provides limited computational power, we decreased the generated_batch_size from 64 to 32, and we also run it for 10 epochs to show our pre-computed training scores. On our local GPU machine, we train the model with generated_batch_size is 64 and run for 200 epochs.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Xq4beSAoQSDA"
   },
   "outputs": [],
   "source": [
    "# training hyperparameters given by code author\n",
    "\n",
    "lr_gen = 0.0001 #Learning rate for generator\n",
    "lr_dis = 0.0001 #Learning rate for discriminator\n",
    "latent_dim = 1024 #Latent dimension\n",
    "gener_batch_size = 32 #Batch size for generator\n",
    "dis_batch_size = 32 #Batch size for discriminator\n",
    "epoch = 10 #Number of epoch\n",
    "weight_decay = 1e-3 #Weight decay\n",
    "drop_rate = 0.5 #dropout\n",
    "n_critic = 5 #\n",
    "max_iter = 500000\n",
    "img_name = \"img_name\"\n",
    "lr_decay = True\n",
    "\n",
    "# architecture details by authors\n",
    "image_size = 32 #H,W size of image for discriminator\n",
    "initial_size = 8 #Initial size for generator\n",
    "patch_size = 4 #Patch size for generated image\n",
    "num_classes = 1 #Number of classes for discriminator \n",
    "output_dir = 'checkpoint' #saved model path\n",
    "dim = 384 #Embedding dimension \n",
    "optimizer = 'Adam' #Optimizer\n",
    "loss = \"wgangp_eps\" #Loss function\n",
    "phi = 1 #\n",
    "beta1 = 0 #\n",
    "beta2 = 0.99 #\n",
    "diff_aug = \"translation,cutout,color\" #data augmentation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I7iL1KYpMyx5"
   },
   "source": [
    "## Training & Saving Model for CIFAR-10\n",
    "As we mentioned above we run the training for 10 epochs due to limitation of Google Colab and showed the decrease in FID score from 253 to 138 in 10 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cS5Oxb7WcAts",
    "outputId": "90caa77c-f93d-48c1-cc46-1e24c5357d09"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/content/drive/My Drive/Colab Notebooks/last_version_28_06/utils.py:64: UserWarning: nn.init.xavier_uniform is now deprecated in favor of nn.init.xavier_uniform_.\n",
      "  nn.init.xavier_uniform(m.weight.data, 1.)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Discriminator(\n",
       "  (patches): ImgPatches(\n",
       "    (patch_embed): Conv2d(3, 384, kernel_size=(4, 4), stride=(4, 4))\n",
       "  )\n",
       "  (droprate): Dropout(p=0.5, inplace=False)\n",
       "  (TransfomerEncoder): TransformerEncoder(\n",
       "    (Encoder_Blocks): ModuleList(\n",
       "      (0): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): Encoder_Block(\n",
       "        (ln1): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (attn): Attention(\n",
       "          (qkv): Linear(in_features=384, out_features=1152, bias=False)\n",
       "          (attention_dropout): Dropout(p=0.5, inplace=False)\n",
       "          (out): Sequential(\n",
       "            (0): Linear(in_features=384, out_features=384, bias=True)\n",
       "            (1): Dropout(p=0.5, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (ln2): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): MLP(\n",
       "          (fc1): Linear(in_features=384, out_features=1536, bias=True)\n",
       "          (act): GELU()\n",
       "          (fc2): Linear(in_features=1536, out_features=384, bias=True)\n",
       "          (droprateout): Dropout(p=0.5, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (norm): LayerNorm((384,), eps=1e-05, elementwise_affine=True)\n",
       "  (out): Linear(in_features=384, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 41,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    dev = \"cuda:0\"\n",
    "else:\n",
    "    dev = \"cpu\"\n",
    "\n",
    "device = torch.device(dev)\n",
    "\n",
    "generator= Generator(depth1=5, depth2=4, depth3=2, initial_size=8, dim=384, heads=4, mlp_ratio=4, drop_rate=0.5)#,device = device)\n",
    "generator.to(device)\n",
    "\n",
    "discriminator = Discriminator(diff_aug = diff_aug, image_size=32, patch_size=4, input_channel=3, num_classes=1,\n",
    "                 dim=384, depth=7, heads=4, mlp_ratio=4,\n",
    "                 drop_rate=0.5)\n",
    "discriminator.to(device)\n",
    "\n",
    "\n",
    "generator.apply(inits_weight)\n",
    "discriminator.apply(inits_weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-FrwqrkneKdO",
    "outputId": "00d3377d-b405-4bec-b98e-659b03dc1c4a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "optimizer: Adam\n"
     ]
    }
   ],
   "source": [
    "if optimizer == 'Adam':\n",
    "    optim_gen = optim.Adam(filter(lambda p: p.requires_grad, generator.parameters()), lr=lr_gen, betas=(beta1, beta2))\n",
    "\n",
    "    optim_dis = optim.Adam(filter(lambda p: p.requires_grad, discriminator.parameters()),lr=lr_dis, betas=(beta1, beta2))\n",
    "elif optimizer == 'SGD':\n",
    "    optim_gen = optim.SGD(filter(lambda p: p.requires_grad, generator.parameters()),\n",
    "                lr=lr_gen, momentum=0.9)\n",
    "\n",
    "    optim_dis = optim.SGD(filter(lambda p: p.requires_grad, discriminator.parameters()),\n",
    "                lr=lr_dis, momentum=0.9)\n",
    "\n",
    "elif optimizer == 'RMSprop':\n",
    "    optim_gen = optim.RMSprop(filter(lambda p: p.requires_grad, discriminator.parameters()),\n",
    "                lr=lr_dis, eps=1e-08, weight_decay=weight_decay, momentum=0, centered=False)\n",
    "\n",
    "    optim_dis = optim.RMSprop(filter(lambda p: p.requires_grad, discriminator.parameters()), lr=lr_dis, eps=1e-08, weight_decay=weight_decay, momentum=0, centered=False)\n",
    "\n",
    "gen_scheduler = LinearLrDecay(optim_gen, lr_gen, 0.0, 0, max_iter * n_critic)\n",
    "dis_scheduler = LinearLrDecay(optim_dis, lr_dis, 0.0, 0, max_iter * n_critic)\n",
    "\n",
    "#RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)\n",
    "\n",
    "print(\"optimizer:\",optimizer)\n",
    "\n",
    "fid_stat = 'fid_stat/fid_stats_cifar10_train.npz'\n",
    "\n",
    "writer=SummaryWriter()\n",
    "writer_dict = {'writer':writer}\n",
    "writer_dict[\"train_global_steps\"]=0\n",
    "writer_dict[\"valid_global_steps\"]=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oqrgmomNf01n"
   },
   "outputs": [],
   "source": [
    "def compute_gradient_penalty(D, real_samples, fake_samples, phi):\n",
    "    \"\"\"Calculates the gradient penalty loss for WGAN GP\"\"\"\n",
    "    # Random weight term for interpolation between real and fake samples\n",
    "    alpha = torch.Tensor(np.random.random((real_samples.size(0), 1, 1, 1))).to(real_samples.get_device())\n",
    "    # Get random interpolation between real and fake samples\n",
    "    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)\n",
    "    d_interpolates = D(interpolates)\n",
    "    fake = torch.ones([real_samples.shape[0], 1], requires_grad=False).to(real_samples.get_device())\n",
    "    # Get gradient w.r.t. interpolates\n",
    "    gradients = torch.autograd.grad(\n",
    "        outputs=d_interpolates,\n",
    "        inputs=interpolates,\n",
    "        grad_outputs=fake,\n",
    "        create_graph=True,\n",
    "        retain_graph=True,\n",
    "        only_inputs=True,\n",
    "    )[0]\n",
    "    gradients = gradients.contiguous().view(gradients.size(0), -1)\n",
    "    gradient_penalty = ((gradients.norm(2, dim=1) - phi) ** 2).mean()\n",
    "    return gradient_penalty\n",
    "\n",
    "\n",
    "def train(noise,generator, discriminator, optim_gen, optim_dis,\n",
    "        epoch, writer, schedulers, img_size=32, latent_dim = latent_dim,\n",
    "        n_critic = n_critic,\n",
    "        gener_batch_size=gener_batch_size, device=\"cuda:0\"):\n",
    "\n",
    "\n",
    "    writer = writer_dict['writer']\n",
    "    gen_step = 0\n",
    "\n",
    "    generator = generator.train()\n",
    "    discriminator = discriminator.train()\n",
    "\n",
    "    transform = transforms.Compose([transforms.Resize(size=(img_size, img_size)),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n",
    "\n",
    "    train_set = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)\n",
    "    train_loader = torch.utils.data.DataLoader(dataset=train_set, batch_size=30, shuffle=True)\n",
    "\n",
    "    for index, (img, _) in enumerate(train_loader):\n",
    "\n",
    "        global_steps = writer_dict['train_global_steps']\n",
    "\n",
    "        real_imgs = img.type(torch.cuda.FloatTensor)\n",
    "\n",
    "        noise = torch.cuda.FloatTensor(np.random.normal(0, 1, (img.shape[0], latent_dim)))#noise(img, latent_dim)#= args.latent_dim)\n",
    "\n",
    "        optim_dis.zero_grad()\n",
    "        real_valid=discriminator(real_imgs)\n",
    "        fake_imgs = generator(noise).detach()\n",
    "        \n",
    "        #assert fake_imgs.size() == real_imgs.size(), f\"fake_imgs.size(): {fake_imgs.size()} real_imgs.size(): {real_imgs.size()}\"\n",
    "\n",
    "        fake_valid = discriminator(fake_imgs)\n",
    "\n",
    "        if loss == 'hinge':\n",
    "            loss_dis = torch.mean(nn.ReLU(inplace=True)(1.0 - real_valid)).to(device) + torch.mean(nn.ReLU(inplace=True)(1 + fake_valid)).to(device)\n",
    "        elif loss == 'wgangp_eps':\n",
    "            gradient_penalty = compute_gradient_penalty(discriminator, real_imgs, fake_imgs.detach(), phi)\n",
    "            loss_dis = -torch.mean(real_valid) + torch.mean(fake_valid) + gradient_penalty * 10 / (phi ** 2)         \n",
    "\n",
    "        loss_dis.backward()\n",
    "        optim_dis.step()\n",
    "\n",
    "        writer.add_scalar(\"loss_dis\", loss_dis.item(), global_steps)\n",
    "\n",
    "        if global_steps % n_critic == 0:\n",
    "\n",
    "            optim_gen.zero_grad()\n",
    "            if schedulers:\n",
    "                gen_scheduler, dis_scheduler = schedulers\n",
    "                g_lr = gen_scheduler.step(global_steps)\n",
    "                d_lr = dis_scheduler.step(global_steps)\n",
    "                writer.add_scalar('LR/g_lr', g_lr, global_steps)\n",
    "                writer.add_scalar('LR/d_lr', d_lr, global_steps)\n",
    "\n",
    "            gener_noise = torch.cuda.FloatTensor(np.random.normal(0, 1, (gener_batch_size, latent_dim)))\n",
    "\n",
    "            generated_imgs= generator(gener_noise)\n",
    "            fake_valid = discriminator(generated_imgs)\n",
    "\n",
    "            gener_loss = -torch.mean(fake_valid).to(device)\n",
    "            gener_loss.backward()\n",
    "            optim_gen.step()\n",
    "            writer.add_scalar(\"gener_loss\", gener_loss.item(), global_steps)\n",
    "\n",
    "            gen_step += 1\n",
    "\n",
    "            #writer_dict['train_global_steps'] = global_steps + 1\n",
    "\n",
    "        if gen_step and index % 100 == 0:\n",
    "            sample_imgs = generated_imgs[:25]\n",
    "            img_grid = make_grid(sample_imgs, nrow=5, normalize=True, scale_each=True)\n",
    "            save_image(sample_imgs, f'generated_images/generated_img_{epoch}_{index % len(train_loader)}.jpg', nrow=5, normalize=True, scale_each=True)            \n",
    "            tqdm.write(\"[Epoch %d] [Batch %d/%d] [D loss: %f] [G loss: %f]\" %\n",
    "                (epoch+1, index % len(train_loader), len(train_loader), loss_dis.item(), gener_loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wA0coJ7kgE5X"
   },
   "outputs": [],
   "source": [
    "def validate(generator, writer_dict, fid_stat):\n",
    "\n",
    "\n",
    "        writer = writer_dict['writer']\n",
    "        global_steps = writer_dict['valid_global_steps']\n",
    "\n",
    "        generator = generator.eval()\n",
    "        fid_score = get_fid(fid_stat, epoch, generator, num_img=5000, val_batch_size=60*2, latent_dim=1024, writer_dict=None, cls_idx=None)\n",
    "\n",
    "\n",
    "        print(f\"FID score: {fid_score}\")\n",
    "\n",
    "        writer.add_scalar('FID_score', fid_score, global_steps)\n",
    "\n",
    "        writer_dict['valid_global_steps'] = global_steps + 1\n",
    "        return fid_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "IvwXcTahgPAV",
    "outputId": "b92e59f3-23d2-42f1-cf40-c7cb1f2e42e4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "[Epoch 1] [Batch 0/1667] [D loss: 2.170649] [G loss: 0.804360]\n",
      "[Epoch 1] [Batch 100/1667] [D loss: -8.304835] [G loss: 10.228638]\n",
      "[Epoch 1] [Batch 200/1667] [D loss: -3.785351] [G loss: 8.178776]\n",
      "[Epoch 1] [Batch 300/1667] [D loss: -1.363279] [G loss: 8.804716]\n",
      "[Epoch 1] [Batch 400/1667] [D loss: 3.654663] [G loss: 4.162188]\n",
      "[Epoch 1] [Batch 500/1667] [D loss: -1.666954] [G loss: 10.869068]\n",
      "[Epoch 1] [Batch 600/1667] [D loss: -0.912084] [G loss: 7.093612]\n",
      "[Epoch 1] [Batch 700/1667] [D loss: 0.081573] [G loss: 6.099335]\n",
      "[Epoch 1] [Batch 800/1667] [D loss: -0.579418] [G loss: 5.591341]\n",
      "[Epoch 1] [Batch 900/1667] [D loss: -0.207472] [G loss: 5.181511]\n",
      "[Epoch 1] [Batch 1000/1667] [D loss: -0.217765] [G loss: 3.403362]\n",
      "[Epoch 1] [Batch 1100/1667] [D loss: -0.506849] [G loss: 4.797401]\n",
      "[Epoch 1] [Batch 1200/1667] [D loss: 0.399069] [G loss: 3.140276]\n",
      "[Epoch 1] [Batch 1300/1667] [D loss: 0.596609] [G loss: 4.660552]\n",
      "[Epoch 1] [Batch 1400/1667] [D loss: -0.246731] [G loss: 4.949975]\n",
      "[Epoch 1] [Batch 1500/1667] [D loss: 0.416937] [G loss: 5.639207]\n",
      "[Epoch 1] [Batch 1600/1667] [D loss: -1.086679] [G loss: 6.080180]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.01it/s]\n",
      "/usr/local/lib/python3.7/dist-packages/torchvision/models/inception.py:82: FutureWarning: The default weight initialization of inception_v3 will be changed in future releases of torchvision. If you wish to keep the old behavior (which leads to long initialization times due to scipy/scipy#11299), please set init_weights=True.\n",
      "  ' due to scipy/scipy#11299), please set init_weights=True.', FutureWarning)\n",
      "  3%|▎         | 3/98 [00:00<00:05, 15.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 253.59642028808594\n",
      "FID score: 253.59642028808594 - best ID score: 10000.0 || @ epoch 1.\n",
      "Saved Latest Model!\n",
      "Files already downloaded and verified\n",
      "[Epoch 2] [Batch 0/1667] [D loss: -1.014748] [G loss: 5.774697]\n",
      "[Epoch 2] [Batch 100/1667] [D loss: -0.355909] [G loss: 6.370594]\n",
      "[Epoch 2] [Batch 200/1667] [D loss: 1.714307] [G loss: 3.597477]\n",
      "[Epoch 2] [Batch 300/1667] [D loss: -0.131275] [G loss: 5.499962]\n",
      "[Epoch 2] [Batch 400/1667] [D loss: 0.279393] [G loss: 4.044170]\n",
      "[Epoch 2] [Batch 500/1667] [D loss: -0.194252] [G loss: 4.526648]\n",
      "[Epoch 2] [Batch 600/1667] [D loss: 0.094030] [G loss: 3.439218]\n",
      "[Epoch 2] [Batch 700/1667] [D loss: 1.523710] [G loss: 4.529868]\n",
      "[Epoch 2] [Batch 800/1667] [D loss: -0.376306] [G loss: 3.141362]\n",
      "[Epoch 2] [Batch 900/1667] [D loss: -1.031662] [G loss: 2.987192]\n",
      "[Epoch 2] [Batch 1000/1667] [D loss: -0.558286] [G loss: 2.625803]\n",
      "[Epoch 2] [Batch 1100/1667] [D loss: 0.100764] [G loss: 3.178997]\n",
      "[Epoch 2] [Batch 1200/1667] [D loss: -0.225897] [G loss: 2.925527]\n",
      "[Epoch 2] [Batch 1300/1667] [D loss: 0.104118] [G loss: 3.194591]\n",
      "[Epoch 2] [Batch 1400/1667] [D loss: -0.338322] [G loss: 3.978561]\n",
      "[Epoch 2] [Batch 1500/1667] [D loss: 0.065034] [G loss: 3.397743]\n",
      "[Epoch 2] [Batch 1600/1667] [D loss: 0.426229] [G loss: 3.736301]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.63it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 256.2900390625\n",
      "FID score: 256.2900390625 - best ID score: 253.59642028808594 || @ epoch 2.\n",
      "Files already downloaded and verified\n",
      "[Epoch 3] [Batch 0/1667] [D loss: -0.372902] [G loss: 1.947552]\n",
      "[Epoch 3] [Batch 100/1667] [D loss: 0.475183] [G loss: 1.971642]\n",
      "[Epoch 3] [Batch 200/1667] [D loss: -0.481462] [G loss: 1.334552]\n",
      "[Epoch 3] [Batch 300/1667] [D loss: 0.233310] [G loss: 1.891104]\n",
      "[Epoch 3] [Batch 400/1667] [D loss: -0.011959] [G loss: 2.165990]\n",
      "[Epoch 3] [Batch 500/1667] [D loss: 0.746799] [G loss: 1.237505]\n",
      "[Epoch 3] [Batch 600/1667] [D loss: 0.324373] [G loss: 1.170042]\n",
      "[Epoch 3] [Batch 700/1667] [D loss: 0.274361] [G loss: 2.603842]\n",
      "[Epoch 3] [Batch 800/1667] [D loss: -0.035530] [G loss: 2.541132]\n",
      "[Epoch 3] [Batch 900/1667] [D loss: 0.114036] [G loss: 2.654665]\n",
      "[Epoch 3] [Batch 1000/1667] [D loss: -0.258933] [G loss: 1.853575]\n",
      "[Epoch 3] [Batch 1100/1667] [D loss: 0.695169] [G loss: 2.649147]\n",
      "[Epoch 3] [Batch 1200/1667] [D loss: -0.254932] [G loss: 2.856803]\n",
      "[Epoch 3] [Batch 1300/1667] [D loss: 0.157940] [G loss: 3.241687]\n",
      "[Epoch 3] [Batch 1400/1667] [D loss: 0.394839] [G loss: 2.448287]\n",
      "[Epoch 3] [Batch 1500/1667] [D loss: 0.150060] [G loss: 2.370821]\n",
      "[Epoch 3] [Batch 1600/1667] [D loss: 0.239766] [G loss: 3.224324]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:05, 17.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:16<00:00,  5.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 207.31369018554688\n",
      "FID score: 207.31369018554688 - best ID score: 253.59642028808594 || @ epoch 3.\n",
      "Files already downloaded and verified\n",
      "[Epoch 4] [Batch 0/1667] [D loss: -0.365915] [G loss: 3.485483]\n",
      "[Epoch 4] [Batch 100/1667] [D loss: 0.443299] [G loss: 3.688633]\n",
      "[Epoch 4] [Batch 200/1667] [D loss: -0.308660] [G loss: 2.951671]\n",
      "[Epoch 4] [Batch 300/1667] [D loss: 0.600909] [G loss: 1.788807]\n",
      "[Epoch 4] [Batch 400/1667] [D loss: 0.937112] [G loss: 1.625080]\n",
      "[Epoch 4] [Batch 500/1667] [D loss: 0.299240] [G loss: 2.000552]\n",
      "[Epoch 4] [Batch 600/1667] [D loss: 0.024430] [G loss: 2.145702]\n",
      "[Epoch 4] [Batch 700/1667] [D loss: -0.568093] [G loss: 1.746835]\n",
      "[Epoch 4] [Batch 800/1667] [D loss: 0.112734] [G loss: 0.780115]\n",
      "[Epoch 4] [Batch 900/1667] [D loss: 0.062658] [G loss: 1.361718]\n",
      "[Epoch 4] [Batch 1000/1667] [D loss: -0.388711] [G loss: 2.048912]\n",
      "[Epoch 4] [Batch 1100/1667] [D loss: 0.292507] [G loss: 2.032061]\n",
      "[Epoch 4] [Batch 1200/1667] [D loss: 0.602682] [G loss: 1.755839]\n",
      "[Epoch 4] [Batch 1300/1667] [D loss: -0.002767] [G loss: 2.913162]\n",
      "[Epoch 4] [Batch 1400/1667] [D loss: 0.379667] [G loss: 2.334970]\n",
      "[Epoch 4] [Batch 1500/1667] [D loss: 0.171352] [G loss: 3.343623]\n",
      "[Epoch 4] [Batch 1600/1667] [D loss: 0.395515] [G loss: 3.462258]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  3.98it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:16<00:00,  5.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 172.43017578125\n",
      "FID score: 172.43017578125 - best ID score: 253.59642028808594 || @ epoch 4.\n",
      "Files already downloaded and verified\n",
      "[Epoch 5] [Batch 0/1667] [D loss: 0.172841] [G loss: 2.592175]\n",
      "[Epoch 5] [Batch 100/1667] [D loss: 0.496616] [G loss: 2.731457]\n",
      "[Epoch 5] [Batch 200/1667] [D loss: 0.140289] [G loss: 2.167178]\n",
      "[Epoch 5] [Batch 300/1667] [D loss: -0.197189] [G loss: 2.397559]\n",
      "[Epoch 5] [Batch 400/1667] [D loss: 0.439036] [G loss: 2.326895]\n",
      "[Epoch 5] [Batch 500/1667] [D loss: 0.547368] [G loss: 1.894225]\n",
      "[Epoch 5] [Batch 600/1667] [D loss: 0.337562] [G loss: 2.094888]\n",
      "[Epoch 5] [Batch 700/1667] [D loss: 0.304808] [G loss: 2.826921]\n",
      "[Epoch 5] [Batch 800/1667] [D loss: 0.845009] [G loss: 1.749743]\n",
      "[Epoch 5] [Batch 900/1667] [D loss: 0.175082] [G loss: 1.266261]\n",
      "[Epoch 5] [Batch 1000/1667] [D loss: 0.731524] [G loss: 0.654840]\n",
      "[Epoch 5] [Batch 1100/1667] [D loss: 0.080363] [G loss: 1.845362]\n",
      "[Epoch 5] [Batch 1200/1667] [D loss: -0.123356] [G loss: 2.430002]\n",
      "[Epoch 5] [Batch 1300/1667] [D loss: -0.121739] [G loss: 2.188072]\n",
      "[Epoch 5] [Batch 1400/1667] [D loss: -0.271561] [G loss: 1.371209]\n",
      "[Epoch 5] [Batch 1500/1667] [D loss: 0.742014] [G loss: 0.496528]\n",
      "[Epoch 5] [Batch 1600/1667] [D loss: 0.103187] [G loss: 1.180739]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 157.23953247070312\n",
      "FID score: 157.23953247070312 - best ID score: 253.59642028808594 || @ epoch 5.\n",
      "Files already downloaded and verified\n",
      "[Epoch 6] [Batch 0/1667] [D loss: 0.371292] [G loss: 0.855299]\n",
      "[Epoch 6] [Batch 100/1667] [D loss: 0.441028] [G loss: 0.614191]\n",
      "[Epoch 6] [Batch 200/1667] [D loss: -0.097045] [G loss: 0.962588]\n",
      "[Epoch 6] [Batch 300/1667] [D loss: 0.367652] [G loss: 0.857501]\n",
      "[Epoch 6] [Batch 400/1667] [D loss: 0.638341] [G loss: 1.082901]\n",
      "[Epoch 6] [Batch 500/1667] [D loss: 0.210048] [G loss: 1.923066]\n",
      "[Epoch 6] [Batch 600/1667] [D loss: -0.016896] [G loss: 1.761053]\n",
      "[Epoch 6] [Batch 700/1667] [D loss: 0.193946] [G loss: 1.685530]\n",
      "[Epoch 6] [Batch 800/1667] [D loss: 1.111356] [G loss: 1.994036]\n",
      "[Epoch 6] [Batch 900/1667] [D loss: -0.004096] [G loss: 2.789292]\n",
      "[Epoch 6] [Batch 1000/1667] [D loss: 0.067194] [G loss: 2.250651]\n",
      "[Epoch 6] [Batch 1100/1667] [D loss: 1.058728] [G loss: 1.593802]\n",
      "[Epoch 6] [Batch 1200/1667] [D loss: -0.007716] [G loss: 1.086996]\n",
      "[Epoch 6] [Batch 1300/1667] [D loss: 0.110086] [G loss: 1.331524]\n",
      "[Epoch 6] [Batch 1400/1667] [D loss: 0.503245] [G loss: 0.986329]\n",
      "[Epoch 6] [Batch 1500/1667] [D loss: 0.088736] [G loss: 1.075401]\n",
      "[Epoch 6] [Batch 1600/1667] [D loss: 0.064434] [G loss: 1.135694]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:05, 16.98it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 143.97250366210938\n",
      "FID score: 143.97250366210938 - best ID score: 253.59642028808594 || @ epoch 6.\n",
      "Files already downloaded and verified\n",
      "[Epoch 7] [Batch 0/1667] [D loss: 0.178964] [G loss: 1.296113]\n",
      "[Epoch 7] [Batch 100/1667] [D loss: 0.032568] [G loss: 2.019550]\n",
      "[Epoch 7] [Batch 200/1667] [D loss: 0.251444] [G loss: 2.284503]\n",
      "[Epoch 7] [Batch 300/1667] [D loss: -0.062520] [G loss: 1.776974]\n",
      "[Epoch 7] [Batch 400/1667] [D loss: 0.211068] [G loss: 2.131124]\n",
      "[Epoch 7] [Batch 500/1667] [D loss: 0.180025] [G loss: 1.905816]\n",
      "[Epoch 7] [Batch 600/1667] [D loss: 0.390855] [G loss: 1.178607]\n",
      "[Epoch 7] [Batch 700/1667] [D loss: 0.813857] [G loss: 1.837076]\n",
      "[Epoch 7] [Batch 800/1667] [D loss: 0.585625] [G loss: 2.185049]\n",
      "[Epoch 7] [Batch 900/1667] [D loss: 0.245140] [G loss: 1.398675]\n",
      "[Epoch 7] [Batch 1000/1667] [D loss: 0.579851] [G loss: 1.662638]\n",
      "[Epoch 7] [Batch 1100/1667] [D loss: -0.435777] [G loss: 2.190656]\n",
      "[Epoch 7] [Batch 1200/1667] [D loss: 0.447797] [G loss: 1.233313]\n",
      "[Epoch 7] [Batch 1300/1667] [D loss: 0.249970] [G loss: 1.034499]\n",
      "[Epoch 7] [Batch 1400/1667] [D loss: -0.068920] [G loss: 0.408645]\n",
      "[Epoch 7] [Batch 1500/1667] [D loss: -0.053509] [G loss: 1.569734]\n",
      "[Epoch 7] [Batch 1600/1667] [D loss: 0.144085] [G loss: 1.404664]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 142.53961181640625\n",
      "FID score: 142.53961181640625 - best ID score: 253.59642028808594 || @ epoch 7.\n",
      "Files already downloaded and verified\n",
      "[Epoch 8] [Batch 0/1667] [D loss: 0.223762] [G loss: 2.058608]\n",
      "[Epoch 8] [Batch 100/1667] [D loss: 0.664551] [G loss: 1.365796]\n",
      "[Epoch 8] [Batch 200/1667] [D loss: -0.030917] [G loss: 0.785733]\n",
      "[Epoch 8] [Batch 300/1667] [D loss: 0.205168] [G loss: 1.533132]\n",
      "[Epoch 8] [Batch 400/1667] [D loss: 0.286476] [G loss: 1.991652]\n",
      "[Epoch 8] [Batch 500/1667] [D loss: 0.340443] [G loss: 2.079747]\n",
      "[Epoch 8] [Batch 600/1667] [D loss: 0.504586] [G loss: 1.812056]\n",
      "[Epoch 8] [Batch 700/1667] [D loss: -0.231664] [G loss: 1.432201]\n",
      "[Epoch 8] [Batch 800/1667] [D loss: 0.073055] [G loss: 2.221403]\n",
      "[Epoch 8] [Batch 900/1667] [D loss: 0.441062] [G loss: 1.701555]\n",
      "[Epoch 8] [Batch 1000/1667] [D loss: 0.664099] [G loss: 2.063335]\n",
      "[Epoch 8] [Batch 1100/1667] [D loss: 0.518637] [G loss: 1.571779]\n",
      "[Epoch 8] [Batch 1200/1667] [D loss: -0.353109] [G loss: 1.516699]\n",
      "[Epoch 8] [Batch 1300/1667] [D loss: 0.508980] [G loss: 1.888465]\n",
      "[Epoch 8] [Batch 1400/1667] [D loss: -0.022976] [G loss: 1.655336]\n",
      "[Epoch 8] [Batch 1500/1667] [D loss: 0.133294] [G loss: 2.531805]\n",
      "[Epoch 8] [Batch 1600/1667] [D loss: -0.128449] [G loss: 1.614341]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 135.49749755859375\n",
      "FID score: 135.49749755859375 - best ID score: 253.59642028808594 || @ epoch 8.\n",
      "Files already downloaded and verified\n",
      "[Epoch 9] [Batch 0/1667] [D loss: 0.644349] [G loss: 1.506196]\n",
      "[Epoch 9] [Batch 100/1667] [D loss: 0.430402] [G loss: 1.700323]\n",
      "[Epoch 9] [Batch 200/1667] [D loss: 0.071626] [G loss: 1.859765]\n",
      "[Epoch 9] [Batch 300/1667] [D loss: 0.712856] [G loss: 1.934990]\n",
      "[Epoch 9] [Batch 400/1667] [D loss: -0.230229] [G loss: 1.745375]\n",
      "[Epoch 9] [Batch 500/1667] [D loss: 0.363142] [G loss: 1.438794]\n",
      "[Epoch 9] [Batch 600/1667] [D loss: 0.296138] [G loss: 1.327496]\n",
      "[Epoch 9] [Batch 700/1667] [D loss: 0.589227] [G loss: 0.971769]\n",
      "[Epoch 9] [Batch 800/1667] [D loss: 0.357308] [G loss: 1.326427]\n",
      "[Epoch 9] [Batch 900/1667] [D loss: 0.602394] [G loss: 2.051815]\n",
      "[Epoch 9] [Batch 1000/1667] [D loss: 0.024476] [G loss: 1.609274]\n",
      "[Epoch 9] [Batch 1100/1667] [D loss: 0.068283] [G loss: 1.800756]\n",
      "[Epoch 9] [Batch 1200/1667] [D loss: -0.241045] [G loss: 1.714448]\n",
      "[Epoch 9] [Batch 1300/1667] [D loss: 0.288398] [G loss: 1.303718]\n",
      "[Epoch 9] [Batch 1400/1667] [D loss: 0.379692] [G loss: 1.425779]\n",
      "[Epoch 9] [Batch 1500/1667] [D loss: 0.074182] [G loss: 1.691209]\n",
      "[Epoch 9] [Batch 1600/1667] [D loss: 0.559807] [G loss: 1.350217]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.00it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 151.22991943359375\n",
      "FID score: 151.22991943359375 - best ID score: 253.59642028808594 || @ epoch 9.\n",
      "Files already downloaded and verified\n",
      "[Epoch 10] [Batch 0/1667] [D loss: 0.345630] [G loss: 0.840272]\n",
      "[Epoch 10] [Batch 100/1667] [D loss: 0.222541] [G loss: 1.538569]\n",
      "[Epoch 10] [Batch 200/1667] [D loss: 0.289131] [G loss: 1.055008]\n",
      "[Epoch 10] [Batch 300/1667] [D loss: 0.186299] [G loss: 1.278053]\n",
      "[Epoch 10] [Batch 400/1667] [D loss: -0.148766] [G loss: 0.599416]\n",
      "[Epoch 10] [Batch 500/1667] [D loss: 0.306913] [G loss: 1.451300]\n",
      "[Epoch 10] [Batch 600/1667] [D loss: -0.133036] [G loss: 1.377939]\n",
      "[Epoch 10] [Batch 700/1667] [D loss: 0.124857] [G loss: 1.209821]\n",
      "[Epoch 10] [Batch 800/1667] [D loss: 0.867689] [G loss: 1.488301]\n",
      "[Epoch 10] [Batch 900/1667] [D loss: 0.344258] [G loss: 1.940454]\n",
      "[Epoch 10] [Batch 1000/1667] [D loss: 0.264276] [G loss: 1.534945]\n",
      "[Epoch 10] [Batch 1100/1667] [D loss: -0.213311] [G loss: 1.180523]\n",
      "[Epoch 10] [Batch 1200/1667] [D loss: 0.196303] [G loss: 1.627282]\n",
      "[Epoch 10] [Batch 1300/1667] [D loss: -0.035001] [G loss: 1.912299]\n",
      "[Epoch 10] [Batch 1400/1667] [D loss: 0.516558] [G loss: 1.200998]\n",
      "[Epoch 10] [Batch 1500/1667] [D loss: -0.138443] [G loss: 1.163946]\n",
      "[Epoch 10] [Batch 1600/1667] [D loss: -0.499637] [G loss: 1.239853]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  4.01it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.56it/s]\n",
      "sample images:   0%|          | 0/41 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 138.50048828125\n",
      "FID score: 138.50048828125 - best ID score: 253.59642028808594 || @ epoch 10.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample images: 100%|██████████| 41/41 [00:10<00:00,  3.88it/s]\n",
      "  3%|▎         | 3/98 [00:00<00:06, 15.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Warning: number of images is not a multiple of the batch size. Some samples are going to be ignored.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 98/98 [00:17<00:00,  5.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FID score: 138.24746704101562\n"
     ]
    }
   ],
   "source": [
    "best = 1e4\n",
    "\n",
    "for epoch in range(epoch):\n",
    "\n",
    "    lr_schedulers = (gen_scheduler, dis_scheduler) if lr_decay else None\n",
    "\n",
    "    train(noise, generator, discriminator, optim_gen, optim_dis,\n",
    "    epoch, writer, lr_schedulers,img_size=32, latent_dim = latent_dim,\n",
    "    n_critic = n_critic,\n",
    "    gener_batch_size=gener_batch_size)\n",
    "\n",
    "    checkpoint = {'epoch':epoch, 'best_fid':best}\n",
    "    checkpoint['generator_state_dict'] = generator.state_dict()\n",
    "    checkpoint['discriminator_state_dict'] = discriminator.state_dict()\n",
    "\n",
    "    score = validate(generator, writer_dict, fid_stat)\n",
    "\n",
    "    print(f'FID score: {score} - best ID score: {best} || @ epoch {epoch+1}.')\n",
    "    if epoch == 0 or epoch > 30:\n",
    "        if score < best:\n",
    "            save_checkpoint(checkpoint, is_best=(score<best), output_dir=output_dir)\n",
    "            print(\"Saved Latest Model!\")\n",
    "            best = score\n",
    "\n",
    "\n",
    "checkpoint = {'epoch':epoch, 'best_fid':best}\n",
    "checkpoint['generator_state_dict'] = generator.state_dict()\n",
    "checkpoint['discriminator_state_dict'] = discriminator.state_dict()\n",
    "score = validate(generator, writer_dict, fid_stat) ####CHECK AGAIN\n",
    "save_checkpoint(checkpoint,is_best=(score<best), output_dir=output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fmPPrYITnQqf"
   },
   "source": [
    "### **Experimental Result Goals vs. Achieved Results**\n",
    "In this project, we aimed to reproduce qualitative results(generating image samples by CIFAR-10 Dataset) and quantitative results in Table 2 and Table 4 of the original paper that shown below.\n",
    "<table>\n",
    "<tr>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/table2.png\" style=\"width: 400px;\"/> </td>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/table4.png\" style=\"width: 400px;\"/> </td>\n",
    "</tr></table>\n",
    "\n",
    "Since we have limited computational resource and time for the training all size of TransGAN model on CIFAR-10 Dataset, we only trained the largest model with data augmentation, TransGAN-XL, for Table 4 results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WpMhUguJaL4o"
   },
   "source": [
    "## Test Model and Results\n",
    "In this section, we loaded pre-trained model and got the following qualitative and quantitative results.\n",
    "\n",
    "### Qualitative Results\n",
    "The following pictures show our generated images at different epoch numbers. \n",
    "<table>\n",
    "<tr>\n",
    "<td style=\"text-align: center\">0 Epoch</td>\n",
    "<td style=\"text-align: center\">40 Epoch</td> \n",
    "<td style=\"text-align: center\">100 Epoch</td>\n",
    "<td style=\"text-align: center\">200 Epoch</td> \n",
    "</tr>\n",
    "<trt>\n",
    "<p align=\"center\"><img width=\"30%\" src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/images/atransgan_cifar.gif\"></p>\n",
    "</tr>\n",
    "<tr>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/0.jpg\" style=\"width: 400px;\"/> </td>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/40.jpg\" style=\"width: 400px;\"/> </td>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/100.jpg\" style=\"width: 400px;\"/> </td>\n",
    "<td> <img src=\"https://raw.githubusercontent.com/asarigun/TransGAN/main/results/200.jpg\" style=\"width: 400px;\"/> </td>\n",
    "</tr>\n",
    "</table>\n",
    "\n",
    "### Quantitative Results\n",
    "As we mentioned above, due to the lack of computational resource, we did our experiments only with the largest model TransGAN-XL and get the following results. We had decided not to implement 'Co-Training with Self-Supervised Auxiliary Task' and 'Locality-Aware Initialization for Self-Attention' since they made only small differences as shown in the paper. The difference between our result and original paper result can be originated in using some different hyperparameters and abovementioned implementation differences. You can see our quantitative result, FID score 26.82, [here](https://github.com/asarigun/TransGAN/blob/main/results/wgangp_eps_optim_Adam_lr_gen_0_0001_lr_dis_0_0001_epoch_200.txt).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "A7FBDk0Xbc4k"
   },
   "source": [
    "## Challenges and Discussions\n",
    "\n",
    "Since the authors did not give detailed hyperparameters for each Transformers Block and Multi-Head Attention Mechanism on version 1, we needed to find the best hyperparameters. Also, in the training part, they did not give detailed hyperparameters such as droprate, weight decay, or batch normalization in version 1. But in the last version of the original paper, authors gave more detailed hyperparameters for training, therefore we got more reasonable results.\n",
    "\n",
    "During the implementation, first we used Hinge loss and faced convergence problem in training. When we tried another loss function, WGAN-GP, that is mentioned in the last version of original paper, we achieved to overcome convergence problem and got better results.\n",
    "\n",
    "As authors didn't share detailed training process in their previous version, we struggled to converge FID score during training. But in the latest version of the original paper, authors provided more details for training so that we achieved to converge FID score in the training.   \n",
    "\n",
    "Due to lack of computational resource, we only trained the largest model, TransGAN-XL in our project. We implemented data augmentation in our model as it is considered crucial for TransGAN in the original paper. We didn't implement 'Co-Training with Self-Supervised Auxiliary Task' and 'Locality-Aware Initialization for Self-Attention' since they make only small differences as shown in the paper."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uC3MAynNh51x"
   },
   "source": [
    "## Citation\n",
    "```\n",
    "@article{jiang2021transgan,\n",
    "  title={TransGAN: Two Transformers Can Make One Strong GAN},\n",
    "  author={Jiang, Yifan and Chang, Shiyu and Wang, Zhangyang},\n",
    "  journal={arXiv preprint arXiv:2102.07074},\n",
    "  year={2021}\n",
    "}\n",
    "```\n",
    "```\n",
    "@article{dosovitskiy2020,\n",
    "  title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},\n",
    "  author={Dosovitskiy, Alexey and Beyer, Lucas and Kolesnikov, Alexander and Weissenborn, Dirk and Zhai, Xiaohua and Unterthiner, Thomas and  Dehghani, Mostafa and Minderer, Matthias and Heigold, Georg and Gelly, Sylvain and Uszkoreit, Jakob and Houlsby, Neil},\n",
    "  journal={arXiv preprint arXiv:2010.11929},\n",
    "  year={2020}\n",
    "}\n",
    "```\n",
    "```\n",
    "@inproceedings{zhao2020diffaugment,\n",
    "  title={Differentiable Augmentation for Data-Efficient GAN Training},\n",
    "  author={Zhao, Shengyu and Liu, Zhijian and Lin, Ji and Zhu, Jun-Yan and Han, Song},\n",
    "  booktitle={Conference on Neural Information Processing Systems (NeurIPS)},\n",
    "  year={2020}\n",
    "}\n",
    "```\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "main_(1).ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "5378bde7bb2c4ef5b231e177b3c7c125": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "5ef80cd84e3f4f4fb320342c0ddbcf74": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_f49c98ac29af4927a55e0998f7fd1b8c",
       "IPY_MODEL_c3b68c830c1743e4ae099886d4610e43"
      ],
      "layout": "IPY_MODEL_c738965f35bd43ed9ca1c627f3ec71e1"
     }
    },
    "757131b61efd4bacb3e3c6befcaaf286": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "bbff1579f370416d8c7116018ad4c0ae": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "c3b68c830c1743e4ae099886d4610e43": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_5378bde7bb2c4ef5b231e177b3c7c125",
      "placeholder": "​",
      "style": "IPY_MODEL_efea6db4dd9a4dd4b680d658051a7dee",
      "value": " 91.2M/91.2M [00:01&lt;00:00, 59.3MB/s]"
     }
    },
    "c738965f35bd43ed9ca1c627f3ec71e1": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "efea6db4dd9a4dd4b680d658051a7dee": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "f49c98ac29af4927a55e0998f7fd1b8c": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_757131b61efd4bacb3e3c6befcaaf286",
      "max": 95628359,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_bbff1579f370416d8c7116018ad4c0ae",
      "value": 95628359
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
